#!/bin/bash

# 安装脚本
echo "开始安装 PhotoID Studio..."

# 检查是否为 root 用户
if [ "$EUID" -ne 0 ]; then 
    echo "请使用 root 权限运行此脚本"
    exit 1
fi

# 1. 安装系统依赖
echo "1. 安装系统依赖..."
apt-get update
apt-get install -y \
    python3 \
    python3-dev \
    python3-venv \
    python3-pip \
    build-essential \
    libgl1-mesa-glx \
    libsm6 \
    libxext6 \
    libxrender1

# 检查 Python 版本
PYTHON_VERSION=$(python3 -c 'import sys; print(".".join(map(str, sys.version_info[:2])))')
echo "检测到 Python 版本: $PYTHON_VERSION"

# 2. 设置工作目录
echo "2. 设置工作目录..."
WORK_DIR="/home/photoid"
mkdir -p $WORK_DIR
cd $WORK_DIR

# 3. 创建虚拟环境
echo "3. 创建 Python 虚拟环境..."
python3 -m venv .venv
source .venv/bin/activate

# 4. 安装 Python 依赖
echo "4. 安装 Python 依赖..."
python3 -m pip install --upgrade pip

# 安装基础依赖
echo "安装基础依赖..."
pip install numpy
pip install pillow

# 安装 OpenCV
echo "安装 OpenCV..."
pip uninstall -y opencv-python-headless opencv-python  # 先卸载可能存在的旧版本
pip install opencv-python  # 安装完整版本而不是 headless 版本

# 安装 rembg 及其依赖
echo "安装 rembg 及其依赖..."
pip install onnxruntime-gpu || pip install onnxruntime
pip install rembg

# 安装其他依赖
echo "安装其他依赖..."
pip install mediapipe
pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu

# 安装InsightFace
echo "安装InsightFace..."
pip install insightface

# 安装Milvus
echo "安装Milvus..."
pip install pymilvus

# 验证安装
echo "验证依赖安装..."
python3 -c "import cv2; print('OpenCV 版本:', cv2.__version__)"
python3 -c "import numpy; print('NumPy 版本:', numpy.__version__)"
python3 -c "import PIL; print('Pillow 版本:', PIL.__version__)"
python3 -c "import mediapipe; print('MediaPipe 版本:', mediapipe.__version__)"
python3 -c "import torch; print('PyTorch 版本:', torch.__version__)"

# 5. 复制项目文件
echo "5. 复制项目文件..."
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
echo "当前脚本目录: $SCRIPT_DIR"

if [ -f "$SCRIPT_DIR/photo_processor.py" ] && [ -f "$SCRIPT_DIR/example.py" ]; then
    echo "找到项目文件，开始复制..."
    cp -r "$SCRIPT_DIR/photo_processor.py" "$WORK_DIR/"
    cp -r "$SCRIPT_DIR/example.py" "$WORK_DIR/"
    echo "文件复制完成"
else
    echo "错误: 未找到必要的项目文件。"
    echo "请确保以下文件存在:"
    echo "- $SCRIPT_DIR/photo_processor.py"
    echo "- $SCRIPT_DIR/example.py"
    exit 1
fi

# 6. 创建命令行链接
echo "6. 创建命令行链接..."
cat > /usr/local/bin/photoid << EOL
#!/bin/bash
cd $WORK_DIR
source $WORK_DIR/.venv/bin/activate
python3 $WORK_DIR/example.py "\$@"
EOL
chmod +x /usr/local/bin/photoid

# 7. 设置权限
echo "7. 设置权限..."
chown -R $SUDO_USER:$SUDO_USER $WORK_DIR

# 下载并启动Milvus (使用docker)
echo "下载并启动Milvus (使用docker)..."
docker pull milvusdb/milvus:latest
docker run -d --name milvus_standalone -p 19530:19530 -p 9091:9091 milvusdb/milvus:latest standalone

echo "安装完成！"
echo "使用方法："
echo "photoid input.jpg [选项]"
echo "运行 photoid --help 查看更多选项" 